#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

import os
from tqdm import tqdm

ARRAY_DICT = {}
ARRAY_DICT["Axiom_CAS-CN1.na36.r1.a1.annot.nutrition.sort.csv"] = "CAS-CN1"
ARRAY_DICT["PharmacoScan_96F.na36.r9.a4.annot.nutrition.sort.csv"] = "PharmacoScan_96F"
ARRAY_DICT["Axiom_PharmacoFocus.na36.r4.a1.annot.nutrition.sort.csv"] = "PharmacoFocus"
ARRAY_DICT["Axiom_PMDA.na36.r7.a8.annot.nutrition.sort.csv"] = "PMDA"
ARRAY_DICT["Axiom_APMRA.na35.r3.a2.annot.nutrition.sort.csv"] = "APMRA"
set_cas_cn1 = set()
set_pharmacoscan_96f = set()
set_pharmacofocus = set()
set_pmda = set()
set_apmra = set()


def load_rsids():
    print("Loading rsids from arrays", end=" ... ")
    load_rsid(set_cas_cn1, "Axiom_CAS-CN1.na36.r1.a1.annot.nutrition.sort.csv")
    load_rsid(set_pharmacoscan_96f, "PharmacoScan_96F.na36.r9.a4.annot.nutrition.sort.csv")
    load_rsid(set_pharmacofocus, "Axiom_PharmacoFocus.na36.r4.a1.annot.nutrition.sort.csv")
    load_rsid(set_pmda, "Axiom_PMDA.na36.r7.a8.annot.nutrition.sort.csv")
    load_rsid(set_apmra, "Axiom_APMRA.na35.r3.a2.annot.nutrition.sort.csv")
    print("complete.")


def load_rsid(array_set: set, filename: str):
    infile = open(os.path.join("/Users/genesis/Seafile/Work/营养知识图谱/nutrition-database-result", filename), "r")
    for line in infile:
        line = line.rstrip("\n")
        if not line.startswith("#") or line.startswith("\"Probe Set ID\""):
            rsid = line.split(",")[2][1:-1]
            if rsid.startswith("rs"):
                array_set.add(rsid)
    infile.close()


def add_array_info(input_path: str, output_path: str):
    infile = open(input_path, "r")
    outfile = open(output_path, "w")
    outfile.write(infile.readline().rstrip("\n") + "\tarray_list\tarray_count\n")
    outfile.flush()
    for line in tqdm(infile):
        line = line.rstrip("\n")
        splitline = line.split("\t")
        rsid = splitline[0]
        rank = int(splitline[6])

        array_list = []
        if rsid in set_cas_cn1:
            array_list.append("CAS-CN1")
        if rsid in set_pharmacoscan_96f:
            array_list.append("PharmacoScan_96F")
        if rsid in set_pharmacofocus:
            array_list.append("PharmacoFocus")
        if rsid in set_pmda:
            array_list.append("PMDA")
        if rsid in set_apmra:
            array_list.append("APMRA")

        array_str = ""
        if len(array_list) > 0:
            array_str = "|".join(array_list)
        array_count = len(array_list)

        if array_count > 0 and rank < 5:
            outfile.write("%s\t%s\t%s\t%s\t%s\t%s\t%d\t%s\t%d\n" % (
                splitline[0], splitline[1], splitline[2], splitline[3], splitline[4], splitline[5], 5, array_str,
                array_count))
        else:
            outfile.write(line + "\t" + array_str + "\t" + str(array_count) + "\n")
        outfile.flush()
    infile.close()
    outfile.close()


if __name__ == '__main__':
    load_rsids()
    print(len(set_cas_cn1))
    print(len(set_pharmacoscan_96f))
    print(len(set_pharmacofocus))
    print(len(set_pmda))
    print(len(set_apmra))
    add_array_info("/Users/genesis/Seafile/Work/营养知识图谱/integrated_list_valid.tsv",
                   "/Users/genesis/Seafile/Work/营养知识图谱/integrated_list_valid_array.tsv")
